from torchvision.datasets import ImageFolder
import torch
import torch.utils.data
from torchvision import transforms
import cv2
from PIL import Image
import json
import os
import random
import shutil

import matplotlib
import matplotlib.pyplot as plt
from PIL import Image
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.optim as optim
import torch.utils.data
import torch.utils.data.distributed
import torch.backends.cudnn
import torchvision.transforms as transforms
from timm.utils import accuracy, AverageMeter
from sklearn.metrics import classification_report
from timm.data.mixup import Mixup
from timm.loss import SoftTargetCrossEntropy
import timm
import cv2
import h5py



# utils
origin_datadir = r'D:/Datasets/NEU-DET/train/images'  # 初始数据集位置
categories = os.listdir(origin_datadir) #['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']
categories_to_id = dict((c, i) for i, c in enumerate(categories)) #{'daisy': 0, 'dandelion': 1, 'roses': 2, 'sunflowers': 3, 'tulips': 4}
id_to_categories = dict((v, k) for k, v in categories_to_id.items()) #{0: 'daisy', 1: 'dandelion', 2: 'roses', 3: 'sunflowers', 4: 'tulips'}

def get_imgs_labels(dir):
    # dir: E:\proj\AI\dataset\build_dataset\datasets\train
    imgs = []
    labels = []
    categories = os.listdir(dir)
    for category in categories:
        category_path = os.path.join(dir, category)
        imgs_path = os.listdir(category_path)
        for img in imgs_path:
            imgs.append(os.path.join(category_path, img))
            labels.append(categories_to_id[category])
    return imgs, labels

class Mydataset(torch.utils.data.Dataset):
    # 类初始化
    def __init__(self, img_paths, labels, transform):
        self.img_paths = img_paths
        self.labels = labels
        self.transform = transform

    # 进行切片
    def __getitem__(self, index): #根据给出的索引进行切片，并对其进行数据处理转换成Tensor，返回成Tensor
        img_path = self.img_paths[index]
        label = self.labels[index]
        # get images from h5file
        f = h5py.File(img_path, 'r')  # 打开h5文件
        # print(f.keys())  # 可以查看所有的主键
        class_tiff = f['class_tiff'][:]
        defect_tiff = f['defect_tiff'][:]
        reference_tiff = f['reference_tiff'][:]
        f.close()


        class_tiff = Image.fromarray(class_tiff)
        defect_tiff = Image.fromarray(defect_tiff)
        reference_tiff = Image.fromarray(reference_tiff)


        if self.transform is not None:
            class_tiff = self.transform(class_tiff)
            defect_tiff = self.transform(defect_tiff)
            reference_tiff = self.transform(reference_tiff)
            # print(f'image size after transform:{image.shape}')#torch.Size([1, 256, 256])
        sample = {'class_tiff': class_tiff, 'defect_tiff': defect_tiff, 'reference_tiff': reference_tiff, 'label': label}
        return sample

        # img_path = self.img_paths[index]
        # label = self.labels[index]
        # # image = Image.open(img_path)
        # image = cv2.imread(img_path)
        # image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        # image = Image.fromarray(image)
        # if self.transform is not None:
        #     image = self.transform(image)
        #     # print(f'image size after transform:{image.shape}')#torch.Size([1, 256, 256])
        # return image, label
    # 返回长度
    def __len__(self):
        return len(self.img_paths)

def get_mean_and_std(train_data):
    train_loader = torch.utils.data.DataLoader(
        train_data, batch_size=1, shuffle=False, num_workers=0,
        pin_memory=True)
    mean = torch.zeros(1)
    std = torch.zeros(1)
    for batch_idx, samples in enumerate(train_loader):
        class_image = samples['defect_tiff']
        for d in range(1):
            mean[d] += class_image[:, d, :, :].mean()
            std[d] += class_image[:, d, :, :].std()
    mean.div_(len(train_data))
    std.div_(len(train_data))
    return list(mean.numpy()), list(std.numpy())

def get_mean_and_std_V2(train_data):
    train_loader = torch.utils.data.DataLoader(
        train_data, batch_size=1, shuffle=False, num_workers=0,
        pin_memory=True)
    mean = torch.zeros(1)
    std = torch.zeros(1)
    for batch_idx, (data, target) in enumerate(train_loader):
        class_image = data
        for d in range(1):#通道维数
            mean[d] += class_image[:, d, :, :].mean()
            std[d] += class_image[:, d, :, :].std()
    mean.div_(len(train_data))
    std.div_(len(train_data))
    return list(mean.numpy()), list(std.numpy())

if __name__ == '__main__':
    train_dataset = ImageFolder(root=r'D:/Datasets/NEU-DET/train/images',
                                transform=transforms.ToTensor())
    print(get_mean_and_std_V2(train_dataset)) # ([0.42368844, 0.42368844, 0.42368844], [0.14975642, 0.14975642, 0.14975642])
    # print(train_dataset)

    # train_imgs, train_labels = get_imgs_labels(r'/home/yeadc/Documents/cxj/datasets_h5/train')
    # dataset_train = Mydataset(train_imgs, train_labels, transform=transforms.ToTensor())
    # print(get_mean_and_std(dataset_train))
    #class: ([0.40272543], [0.13901867])
    #defect:([0.402817], [0.14428316])
    #reference:([0.4062375], [0.14477405])

